import streamlit as st

import argparse
import shutil
import subprocess
import numpy as np
import pandas as pd
from z3 import *
import random
import os
from openai import OpenAI

import random
import inspect

#---------------------约束求解部分--------------------------------------

def read_python(file_name):
    with open(file_name, 'r') as file:
        line = file.readline()
        line = line.split('def')[1]
        line = line.strip()
        content = line
    return content

def read_generate_python(file_name):
    with open('../resources/'+str(file_name), 'r') as file:
        line = file.read()
        content = '"""' + '\n' + line + '\n' + '"""'
    return content

def read_python_file(file_name, content):
    with open(file_name, 'r') as file:
        for line in file:
            if 'assert' in line:
                line = line.split('.')[1]
                line = line.strip()
                content += '\n' + line
    return content

def language_to_python(file_name):

    s1 = read_generate_python(str(file_name))

    s2 = """
    如果这段程序是python语言，直接返回原文，若不是python语言，将这段程序转为python语言，在函数参数后加上':'和数据类型（只返回代码即可）
    """

    question = s1 + s2

    client = OpenAI(
        api_key="sk-067ca2043e80470ba2313779960367f8",
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )
    completion = client.chat.completions.create(
        model="qwen-plus",
        # 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
        messages=[
            {'role': 'system', 'content': 'You are a helpful assistant.'},
            {'role': 'user', 'content': question}],
    )

    # 解析返回的JSON数据并提取content字段
    response_content = completion.choices[0].message.content
    response_content = response_content.split('python')[1]
    response_content = response_content.split('```')[0]
    response_content = response_content.strip()

    with open('../resources/test.py', 'w') as file:
        file.write(response_content)

def generate(file_name):
    # BAT脚本路径
    bat_path = 'test.bat'

    # 运行BAT脚本
    cmd = subprocess.Popen([bat_path, file_name], shell=True)
    cmd.wait()

def display(file_name1):
    file_name = 'test_' + file_name1
    with open('result.txt', 'w') as file:
        content0 = read_python('../resources/'+file_name1)
        content = read_python_file(file_name, content0)
        print(content)
        file.write(content)
    with open('result.txt', 'r') as f:
        pos_case = f.read()
        print(pos_case)
        st.code(pos_case)


#--------------------约束求解部分结束----------------------------------

#--------------------遗传算法部分--------------------------------------
# 遗传算法参数
population_size = 50
num_generations = 100
mutation_rate = 0.1

def generate_initial_population():
    return [{'x': random.randint(1, 15), 'y': random.randint(1, 15)} for _ in range(population_size)]

def fitness(test_case):
    # 计算适应度函数，这里简单地返回测试用例是否满足约束条件
    return int(test_case['x'] + test_case['y'] >= 10 and test_case['x'] <= test_case['y'])

def crossover(parent1, parent2):
    child = {}
    for key in parent1.keys():
        child[key] = parent1[key] if random.random() > 0.5 else parent2[key]
    return child

def mutate(test_case):
    if random.random() < mutation_rate:
        key = random.choice(list(test_case.keys()))
        test_case[key] += random.choice([-1, 1])
    return test_case

def genetic_algorithm():
    global best_individual
    population = generate_initial_population()

    for generation in range(num_generations):
        # 选择
        population.sort(key=lambda tc: fitness(tc), reverse=True)

        # 交叉
        new_population = []
        for i in range(population_size // 2):
            parent1, parent2 = random.sample(population[:population_size // 2], 2)
            child = crossover(parent1, parent2)
            new_population.append(child)
            new_population.append(mutate(child))

        # 替换旧种群
        population = new_population

        # 输出当前世代的最佳个体
        best_individual = max(population, key=fitness)
        print(f"Generation {generation}: Best individual {best_individual} with fitness {fitness(best_individual)}")

    return best_individual

def genetic_generate(s1):

    s2 = """
    提取这段代码的约束条件,'&&'用','代替,用python中的布尔型数组存储每个约束条件,保留变量名（只返回代码即可）
    """

    question = s1 + s2

    client = OpenAI(
        api_key="sk-067ca2043e80470ba2313779960367f8",
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )
    completion = client.chat.completions.create(
        model="qwen-plus",
        # 模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
        messages=[
            {'role': 'system', 'content': 'You are a helpful assistant.'},
            {'role': 'user', 'content': question}],
    )

    # 解析返回的JSON数据并提取content字段
    response_content = completion.choices[0].message.content
    response_content = response_content.split('python')[1]
    response_content = response_content.split('```')[0]
    response_content = response_content.strip()
    print(response_content)
    # 定义变量
    x = Int('x')
    y = Int('y')
    # 定义约束条件
    exec(response_content)

    print(constraints)

    # 使用Z3求解器验证约束条件
    s = Solver()
    s.add(constraints)
    is_sat = s.check()

    # 如果存在解，则获取模型
    solution = None
    if is_sat == sat:
        m = s.model()
        solution = {
            m[x].decl().name(): m[x].as_long(),
            m[y].decl().name(): m[y].as_long()
        }

    # 运行遗传算法
    best_test_case = genetic_algorithm()
    print(f"Best test case found: {best_test_case}")
    return best_test_case

#-------------------遗传算法部分结束--------------------------------

#-------------------网页部分---------------------------------
# 初始化Streamlit界面并设置页面布局为宽屏模式
st.set_page_config(layout="wide")
# 设置页面标题
st.title("测试用例自动生成")

# st.subheader("", divider="rainbow")

tab1, tab2, tab3 = st.tabs(["在线输入", "文件上传", "数据统计"])

with tab1:
    # 添加文本
    txt = st.text_area("", value='', height=300, placeholder="输入程序")

    # 创建3个等宽的列
    col1, col2, col3 = st.columns(3)

    # 在中间列中添加按钮
    with col2:
        first_button, second_button = st.columns(2)

    with open('../resources/' + str('test.py'), 'w') as f0:
        f0.write(txt)

    if first_button.button('用约束算法生成'):
        with st.spinner('正在生成...'):
            language_to_python('test.py')
            generate('../resources/test.py')
        st.success('生成完成!')
        display('test.py')

    if second_button.button('用遗传算法生成'):
        with st.spinner('正在生成...'):
            result = genetic_generate(txt)
        st.success('生成完成!')
        st.code(result)
with tab2:
    uploaded_files = st.file_uploader("选择文件", accept_multiple_files=True)
    first_button, second_button = st.columns(2)
    for uploaded_file in uploaded_files:
        if first_button.button(uploaded_file.name + '用约束算法生成'):
            language_to_python(uploaded_file.name)
            with st.spinner('正在生成...'):
                generate('../resources/test.py')
            st.success('生成完成!')
            display('test.py')
        if second_button.button(uploaded_file.name + '用遗传算法生成'):
            with st.spinner('正在生成...'):
                with open('../resources/'+str(uploaded_file.name), 'r') as f:
                    s1 = f.read()
                    result = genetic_generate(s1)
            st.success('生成完成!')
            st.code(result)

with tab3:
    con_time = [2.605, 2.698, 2.671, 2.787, 2.811,
            2.867, 3.145, 2.975, 2.842, 2.768,
            2.713, 3.219, 2.984, 2.856, 2.950,
            3.023, 2.832, 2.880, 2.920, 2.997]
    genetic_time = [0.747, 0.69081, 0.77094, 0.70976, 0.67344, 0.66847, 0.73138, 0.67589, 0.72822, 0.69765,
           0.67335, 0.71873, 0.77235, 0.75911, 0.75829, 0.69377, 0.88354, 0.72119, 0.89954, 0.99954]


    # 创建DataFrame
    data = {'con_time': con_time,'genetic_time': genetic_time }
    df = pd.DataFrame(data)
    chart_data = pd.DataFrame(df)

    st.write("生成用例数量和生成时间之间的关系图")
    st.line_chart(chart_data)

    chart_data1 = pd.DataFrame(np.random.randn(20, 2), columns=["a", "b"])
    st.bar_chart(chart_data1)

    chart_data2 = pd.DataFrame(np.random.randn(20, 2), columns=["a", "b"])
    st.area_chart(chart_data2)
